#include <iostream>
#include <vector>
#include <string>
#include <sstream>
#include <fstream>
#include <algorithm>
#include <iomanip>
#include <cstdint>
#include <omp.h>
using namespace std;

typedef vector<uint64_t> BitVector;

string trim(const string& str) {
    size_t first = str.find_first_not_of(" \t\r\n");
    if (first == string::npos)
        return "";
    size_t last = str.find_last_not_of(" \t\r\n");
    return str.substr(first, (last - first + 1));
}

vector<int> splitClause(const string& line) {
    vector<int> literals;
    stringstream ss(line);
    int num;
    while (ss >> num) {
        if (num == 0)
            break;
        literals.push_back(num);
    }
    return literals;
}

BitVector createAllOnesBitVector(size_t size) {
    size_t num_blocks = (size + 63) / 64;
    BitVector bv(num_blocks, ~0ULL);
    if (size % 64 != 0) {
        bv.back() = (1ULL << (size % 64)) - 1;
    }
    return bv;
}

inline void clearBit(BitVector& bv, uint64_t index) {
    const size_t block = index >> 6;
    const size_t offset = index & 0x3F;
    bv[block] &= ~(1ULL << offset);
}

inline bool checkBit(const BitVector& bv, uint64_t index) {
    const size_t block = index >> 6;
    const size_t offset = index & 0x3F;
    return (bv[block] >> offset) & 1ULL;
}

void processExternalAssignmentsOpt(const vector<int>& external_vars, 
                                uint64_t fixed_mask,
                                BitVector& array_bitset,
                                uint64_t size,
                                const vector<uint64_t>& var_mask_map) {
    const int m = external_vars.size();
    const uint64_t total = 1ULL << m;
    
    for (uint64_t i = 0; i < total; ++i) {
        uint64_t variable_mask = 0;
        uint64_t tmp = i;
        int idx = 0;
        
        while (tmp) {
            const int tz = __builtin_ctzll(tmp);
            idx += tz;
            variable_mask |= var_mask_map[external_vars[idx] - 1];
            tmp >>= (tz + 1);
            idx++;
        }
        
        const uint64_t combined = fixed_mask | variable_mask;
        if (combined < size) {
            clearBit(array_bitset, combined);
        }
    }
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        cerr << "Usage: " << argv[0] << " <cnf_file>" << endl;
        return 1;
    }

    cout << "Processing input file: " << argv[1] << endl;
    double start_time = omp_get_wtime();

    ifstream infile(argv[1]);
    if (!infile.is_open()) {
        cerr << "Cannot open file " << argv[1] << endl;
        return 1;
    }

    string line;
    int num_vars = 0;
    int num_clauses = 0;
    vector<vector<int>> clauses;
    bool problem_line_found = false;

    while (getline(infile, line)) {
        string trimmed_line = trim(line);
        if (trimmed_line.empty()) continue;
        if (trimmed_line[0] == 'c') continue;

        if (trimmed_line[0] == 'p') {
            string tmp1, tmp2;
            stringstream ss(trimmed_line);
            ss >> tmp1 >> tmp2 >> num_vars >> num_clauses;
            if (tmp1 != "p" || tmp2 != "cnf") {
                cerr << "Error: Invalid problem line format." << endl;
                return 1;
            }
            problem_line_found = true;
            continue;
        }

        if (!problem_line_found) {
            cerr << "Error: Problem line not found before clauses." << endl;
            return 1;
        }

        vector<int> clause = splitClause(trimmed_line);
        if (!clause.empty()) {
            clauses.push_back(clause);
        }
    }

    infile.close();

    if (!problem_line_found) {
        cerr << "Error: No problem line (p cnf ...) found." << endl;
        return 1;
    }

    if (clauses.size() != num_clauses) {
        cerr << "Warning: Number of clauses read (" << clauses.size() 
             << ") differs from expected (" << num_clauses << ")" << endl;
    }

    cout << "Number of variables: " << num_vars << endl;
    cout << "Number of clauses: " << clauses.size() << endl;

    if (num_vars > 28) {
        cerr << "Warning: High number of variables (" << num_vars 
             << "), performance may degrade." << endl;
    }

    // 预计算变量掩码映射表
    vector<uint64_t> var_mask_map(num_vars);
    for (int i = 0; i < num_vars; ++i) {
        var_mask_map[i] = 1ULL << i;
    }

    const uint64_t total_assignments = 1ULL << num_vars;
    BitVector total_bv = createAllOnesBitVector(total_assignments);

    // 优化1：增加任务粒度（每个任务处理更多子句）
    const int MIN_CLAUSES_PER_TASK = 6; // 根据系统调整
    const int num_tasks = max(1, min((int)clauses.size() / MIN_CLAUSES_PER_TASK, 
                                    omp_get_max_threads() * 4));

    cout << "Using " << num_tasks << " parallel tasks" << endl;

    // 优化2：使用任务私有位向量 + 归约合并
    vector<BitVector> task_bvs(num_tasks);
    
    #pragma omp parallel for schedule(static)
    for (int task_id = 0; task_id < num_tasks; task_id++) {
        // 为每个任务创建私有位向量
        task_bvs[task_id] = createAllOnesBitVector(total_assignments);
        
        // 计算当前任务处理的子句范围
        int start_idx = task_id * clauses.size() / num_tasks;
        int end_idx = (task_id + 1) * clauses.size() / num_tasks;
        
        for (int idx = start_idx; idx < end_idx; idx++) {
            const vector<int>& clause = clauses[idx];
            uint64_t fixed_mask = 0;
            vector<int> clause_vars;
            
            for (int lit : clause) {
                int var = abs(lit);
                clause_vars.push_back(var);
                if (lit < 0) {
                    fixed_mask |= var_mask_map[var - 1];
                }
            }
            
            vector<int> external_vars;
            for (int var = 1; var <= num_vars; var++) {
                if (find(clause_vars.begin(), clause_vars.end(), var) == clause_vars.end()) {
                    external_vars.push_back(var);
                }
            }
            
            processExternalAssignmentsOpt(external_vars, fixed_mask, 
                                         task_bvs[task_id], total_assignments, var_mask_map);
        }
    }

    // 归并所有任务的结果
    for (int i = 0; i < num_tasks; i++) {
        for (size_t j = 0; j < total_bv.size(); j++) {
            total_bv[j] &= task_bvs[i][j];
        }
    }

    // 检索满足所有子句的赋值
    cout << "Searching for satisfying assignments..." << endl;
    vector<uint64_t> sat_assignments;
    
    // 优化3：并行化结果检索
    vector<vector<uint64_t>> private_assignments(omp_get_max_threads());
    
    #pragma omp parallel
    {
        int tid = omp_get_thread_num();
        #pragma omp for schedule(static)
        for (uint64_t assignment = 0; assignment < total_assignments; ++assignment) {
            if (checkBit(total_bv, assignment)) {
                private_assignments[tid].push_back(assignment);
            }
        }
    }
    
    // 合并结果
    for (auto& vec : private_assignments) {
        sat_assignments.insert(sat_assignments.end(), vec.begin(), vec.end());
    }

    // 输出结果
    cout << "SATISFIABLE (" << sat_assignments.size() << " solutions)" << endl;
    
    double end_time = omp_get_wtime();
    cout << "Total time: " << fixed << setprecision(6) << (end_time - start_time) << " seconds" << endl;

    return 0;
}